#include "stdafx.h"

#include "TraceServer.h"
#include <unordered_map>
#include <map>
#include "TraceCollection.h"
#include "TraceSymbol.h"
using namespace std;
using namespace std::tr1;
TraceServer *TraceServer::s_instance = NULL;
namespace {
	//unordered_map<void*, LogMessage> memtraceMap;
	TraceCollection traceCollection;
}
TraceServer* TraceServer::GetInstance()
{
	if (s_instance == NULL){
		s_instance = new TraceServer();
	}
	return s_instance;
}
void TraceServer::Destory(){
	delete s_instance;
}
BOOL TraceServer::Init()
{
	InitializeCriticalSection(&m_rawMessageQueueLock);
	//Create IoCP and pipe
	m_hCompletionPort = CreateIoCompletionPort(INVALID_HANDLE_VALUE, NULL, NULL, 0);
	if (m_hCompletionPort == NULL){
		StatusUpdate(_T("create IO completion Port fail."));
		return FALSE;
	}
	CreatePipe();
	CreateWorkerThread();
	return TRUE;
}

BOOL TraceServer::Attach(DWORD PID)
{
	HANDLE hProcess = 
		OpenProcess(
		PROCESS_CREATE_THREAD | PROCESS_QUERY_INFORMATION | PROCESS_VM_OPERATION | PROCESS_VM_WRITE | PROCESS_VM_READ,
		FALSE, PID);
	if (hProcess == NULL)
	{
		StatusUpdate(_T("invalid PID"));
		return FALSE;
	}
	//get the path of dll to inject, it should be in the same directory of exe
	char   szLibPath [_MAX_PATH]={};
	GetModuleFileNameA(NULL, szLibPath,sizeof(szLibPath));
	char *p =strrchr(szLibPath,'\\');
	if(p == NULL){
		p = szLibPath;
	}else{
		p++;
	}
	strncpy(p,"tracedll.dll",sizeof(szLibPath)- (p - szLibPath)-1);
	szLibPath[_MAX_PATH-1]=0;
	// the address (in the remote process) where szLibPath will be copied to;
	void*  pLibRemote = 0;	
	// base adress of loaded module (==HMODULE);
	DWORD  hLibModule = 0;	

	HMODULE hKernel32 = ::GetModuleHandle(_T("Kernel32"));
	pLibRemote = ::VirtualAllocEx( hProcess, NULL, sizeof(szLibPath), MEM_COMMIT, PAGE_READWRITE );
	if( pLibRemote == NULL )
		return FALSE;
	::WriteProcessMemory(hProcess, pLibRemote, (void*)szLibPath,sizeof(szLibPath),NULL);
	HANDLE hThread = ::CreateRemoteThread( hProcess, NULL, 0,	
		(LPTHREAD_START_ROUTINE) ::GetProcAddress(hKernel32,"LoadLibraryA"), 
		pLibRemote, 0, NULL );
	if( hThread == NULL )
	{
		StatusUpdate(_T("remote thread error"));
		return FALSE;
	}else{
		::WaitForSingleObject( hThread, INFINITE );
		CloseHandle(hThread);
	}
	::VirtualFreeEx(hProcess,pLibRemote,sizeof(szLibPath), MEM_DECOMMIT);
	CloseHandle(hProcess);
	return TRUE;
}

DWORD WINAPI TraceServer::WorkerThread(LPVOID p)
{
	TraceServer *this2 = (TraceServer *)p;
	return this2->WorkerThreadFunc();
}

DWORD TraceServer::WorkerThreadFunc()
{
	BOOL ret;
	DWORD nBytes;
	Client *pClient = NULL;
	LPOVERLAPPED lpo;

	while(TRUE){
		ret = GetQueuedCompletionStatus(m_hCompletionPort, 
			&nBytes, (PULONG_PTR)&pClient, &lpo, INFINITE);

		if (!ret){
			if(lpo == NULL){ 
				//IOCP closed, thread should exit.
				break;
			}
			if ((pClient != NULL) && (GetLastError() == ERROR_BROKEN_PIPE)){
				CString msg;
				msg.Format(_T("client disconnected\n"));	
				CloseConnection(pClient);
				continue;
			}
		}
		if (pClient->fAwaitingAccept) {
			pClient->fAwaitingAccept = FALSE;
			CString msg;
			msg.Format(_T("New client connected\n"));			
			StatusUpdate(msg);
			ret = ReadFile(pClient->hPipe,
				&pClient->messageBuf,
				sizeof(pClient->messageBuf),
				&nBytes,
				pClient);
			CreatePipe();
		}else{
			//get message from client
			unsigned int i = 0;
			EnterCriticalSection(&m_rawMessageQueueLock);
			for (; (nBytes>0) && (i < BUF_SIZE); nBytes-=sizeof(LogMessage)){
				LogMessage &message = pClient->messageBuf[i];
				traceCollection.AddMessage(message);
				++i;					
			}
			LeaveCriticalSection(&m_rawMessageQueueLock);
			ret = ReadFile(pClient->hPipe,
				&pClient->messageBuf,
				sizeof(pClient->messageBuf),
				&nBytes,
				pClient);
		}

	}
	return 0;
}

BOOL TraceServer::CreatePipe()
{
	HANDLE hPipe = CreateNamedPipe(TRACE_PIPE_NAME,         // pipe name
		PIPE_ACCESS_INBOUND |       
		FILE_FLAG_OVERLAPPED,       
		PIPE_TYPE_MESSAGE |         
		PIPE_READMODE_MESSAGE |     
		PIPE_WAIT,                  
		PIPE_UNLIMITED_INSTANCES,   
		0,                          // output buffer size
		0,                          // input buffer size
		20000,                      // client time-out
		NULL);                      // no security attributes
	if (hPipe == INVALID_HANDLE_VALUE) {
		StatusUpdate(_T("CreatePipe fail."));
		return FALSE;
	}
	Client *pClient = new Client();
	ZeroMemory(pClient, sizeof(Client));
	pClient->hPipe = hPipe;
	pClient->fAwaitingAccept = TRUE;
	if (!CreateIoCompletionPort(pClient->hPipe,m_hCompletionPort, (ULONG_PTR)pClient,0)){
		StatusUpdate(_T("Fail to add pipe to IOCP."));
		return FALSE;
	}
	if (!ConnectNamedPipe(hPipe, pClient)) {
		if (GetLastError() != ERROR_IO_PENDING &&
			GetLastError() != ERROR_PIPE_LISTENING) {
				StatusUpdate(_T("ConnectNamedPipe"));
		}
	}
	return TRUE;
}

BOOL TraceServer::CloseConnection(Client *pClient)
{
	if (pClient != NULL){
		if (pClient->hPipe != INVALID_HANDLE_VALUE){
			FlushFileBuffers(pClient->hPipe);
			DisconnectNamedPipe(pClient->hPipe);
			CloseHandle(pClient->hPipe);
		}
		delete pClient;
	}
	return TRUE;
}

BOOL TraceServer::CreateWorkerThread()
{
	DWORD dwThread;
	SYSTEM_INFO SystemInfo;

	GetSystemInfo(&SystemInfo);
	DWORD i =0;
	for ( i = 0; i < SystemInfo.dwNumberOfProcessors; i++) {
		HANDLE hThread =CRTThreadTraits::CreateThread(NULL, 0, WorkerThread, this, 0, &dwThread);
		CloseHandle(hThread);
	}
	return TRUE;
}

BOOL TraceServer::SaveLog(FILE *fp)
{
	CString log;
	CString msg;
	multimap<int, CString> msgMap;

	EnterCriticalSection(&m_rawMessageQueueLock);
	auto groupedMessages = traceCollection.GetGroupedMessageMap(); 
	LeaveCriticalSection(&m_rawMessageQueueLock);
	
	TraceSymbol symbol;
	if(groupedMessages.empty()){
		return TRUE;
	}
	DWORD pid = groupedMessages.begin()->second[0].nProcessId;
	if (!symbol.initSymbol(pid)){
		CString errorMessage;
		errorMessage.Format(_T("symbol initialize failed for process:%d."), pid);
		StatusUpdate(errorMessage);
	}
	auto it = groupedMessages.begin(); 
	
	while(it != groupedMessages.end()){
		auto & allocationLogs = it->second;
		int allocCount = allocationLogs.size();
		int totalSize = 0;
		for (auto ite = allocationLogs.begin(); ite!=allocationLogs.end(); ++ite){
			totalSize += ite->size;
		}
		msg.Format(_T("Allocation count: %d, total size: %d\n"), allocCount, totalSize);
		msg.AppendFormat(_T("Process ID: %d,Call stack:\n"), pid);
		const DWORD_PTR* pStack = it->first.stack;
		for(int i=0; i<MEMORY_STACKS; i++) {
			if(pStack[i] == 0) {
				break;
			}
			msg.AppendFormat(_T("[%d]:0x%p, "),i,pStack[i]);
			CString symbolInfo;
			symbol.GetSymbolNamefromAddress(pStack[i],symbolInfo);
			msg.Append(symbolInfo);
			msg.Append(_T("\n"));
		}
		msgMap.insert(make_pair(totalSize, msg));
		++it;
	}
	//sort the message by totalsize, from big to small
	for( auto msgIte = msgMap.rbegin(); msgIte != msgMap.rend(); ++msgIte){
		log.Append(msgIte->second);
	}
	fprintf(fp,"%s\n",CT2A(log));
	return TRUE;
}

void TraceServer::StatusUpdate(const TCHAR *message)
{
	OutputDebugString(message);
	if (m_statusLogger != NULL){
		m_statusLogger->StatusUpdate(message);
	}
}